/*
 * Copyright 2025 the original author or authors.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      https://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package com.alibaba.cloud.ai.lynxe.subplan.model.vo;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.ai.chat.model.ToolContext;

import com.alibaba.cloud.ai.lynxe.planning.model.vo.PlanTemplateConfigVO;
import com.alibaba.cloud.ai.lynxe.planning.service.IPlanParameterMappingService;
import com.alibaba.cloud.ai.lynxe.planning.service.PlanTemplateService;
import com.alibaba.cloud.ai.lynxe.runtime.entity.vo.PlanExecutionResult;
import com.alibaba.cloud.ai.lynxe.runtime.entity.vo.PlanInterface;
import com.alibaba.cloud.ai.lynxe.runtime.entity.vo.RequestSource;
import com.alibaba.cloud.ai.lynxe.runtime.service.PlanIdDispatcher;
import com.alibaba.cloud.ai.lynxe.runtime.service.PlanningCoordinator;
import com.alibaba.cloud.ai.lynxe.tool.AbstractBaseTool;
import com.alibaba.cloud.ai.lynxe.tool.AsyncToolCallBiFunctionDef;
import com.alibaba.cloud.ai.lynxe.tool.code.ToolExecuteResult;
import com.fasterxml.jackson.databind.ObjectMapper;

/**
 * Wrapper class that extends AbstractBaseTool for FuncAgentToolEntity
 *
 * This allows integration with the existing tool registry system. Implements
 * AsyncToolCallBiFunctionDef to support non-blocking execution in parallel scenarios.
 */
public class SubplanToolWrapper extends AbstractBaseTool<Map<String, Object>>
		implements AsyncToolCallBiFunctionDef<Map<String, Object>> {

	public static final String PARENT_PLAN_ID_ARG_NAME = "PLAN_PARENT_ID_ARG_NAME";

	private static final Logger logger = LoggerFactory.getLogger(SubplanToolWrapper.class);

	private final PlanTemplateConfigVO coordinatorToolConfig;

	private final PlanTemplateConfigVO planTemplateConfig;

	private final String currentPlanId;

	private final String rootPlanId;

	private final PlanTemplateService planTemplateService;

	private final PlanningCoordinator planningCoordinator;

	private final PlanIdDispatcher planIdDispatcher;

	private final ObjectMapper objectMapper;

	private final IPlanParameterMappingService parameterMappingService;

	public SubplanToolWrapper(PlanTemplateConfigVO coordinatorToolConfig, PlanTemplateConfigVO planTemplateConfig,
			String currentPlanId, String rootPlanId, PlanTemplateService planTemplateService,
			PlanningCoordinator planningCoordinator, PlanIdDispatcher planIdDispatcher, ObjectMapper objectMapper,
			IPlanParameterMappingService parameterMappingService) {
		this.coordinatorToolConfig = coordinatorToolConfig;
		this.planTemplateConfig = planTemplateConfig;
		this.currentPlanId = currentPlanId;
		this.rootPlanId = rootPlanId;
		this.planTemplateService = planTemplateService;
		this.planningCoordinator = planningCoordinator;
		this.planIdDispatcher = planIdDispatcher;
		this.objectMapper = objectMapper;
		this.parameterMappingService = parameterMappingService;
	}

	@Override
	public String getServiceGroup() {
		// Get serviceGroup from PlanTemplateConfigVO
		if (planTemplateConfig != null && planTemplateConfig.getServiceGroup() != null
				&& !planTemplateConfig.getServiceGroup().trim().isEmpty()) {
			return planTemplateConfig.getServiceGroup();
		}
		// Fallback to default value if serviceGroup is null/empty
		return "coordinator-tools";
	}

	@Override
	public String getName() {
		// Use PlanTemplateConfigVO title as tool name
		if (planTemplateConfig != null && planTemplateConfig.getTitle() != null
				&& !planTemplateConfig.getTitle().trim().isEmpty()) {
			return planTemplateConfig.getTitle();
		}
		// Fallback to planTemplateId if title is not available
		return coordinatorToolConfig != null && coordinatorToolConfig.getPlanTemplateId() != null
				? coordinatorToolConfig.getPlanTemplateId() : "Unknown Tool";
	}

	@Override
	public String getDescription() {
		if (coordinatorToolConfig != null && coordinatorToolConfig.getToolConfig() != null) {
			return coordinatorToolConfig.getToolConfig().getToolDescription();
		}
		return "";
	}

	@Override
	public String getParameters() {
		// Convert PlanTemplateConfigVO inputSchema to JSON schema format
		try {
			List<PlanTemplateConfigVO.InputSchemaParam> inputSchemaParams = null;
			if (coordinatorToolConfig != null && coordinatorToolConfig.getToolConfig() != null) {
				inputSchemaParams = coordinatorToolConfig.getToolConfig().getInputSchema();
			}

			if (inputSchemaParams == null || inputSchemaParams.isEmpty()) {
				return "{\"type\":\"object\",\"properties\":{},\"additionalProperties\":false}";
			}

			// Convert list of InputSchemaParam to JSON schema format
			Map<String, Object> schema = new HashMap<>();
			schema.put("type", "object");

			Map<String, Object> properties = new HashMap<>();
			List<String> required = new ArrayList<>();

			for (PlanTemplateConfigVO.InputSchemaParam param : inputSchemaParams) {
				if (param == null || param.getName() == null || param.getName().trim().isEmpty()) {
					continue;
				}

				Map<String, Object> paramSchema = new HashMap<>();
				String paramType = param.getType() != null ? param.getType().toLowerCase() : "string";
				paramSchema.put("type", paramType);

				if (param.getDescription() != null && !param.getDescription().trim().isEmpty()) {
					paramSchema.put("description", param.getDescription());
				}

				properties.put(param.getName(), paramSchema);

				if (param.getRequired() != null && param.getRequired()) {
					required.add(param.getName());
				}
			}

			schema.put("properties", properties);
			if (!required.isEmpty()) {
				schema.put("required", required);
			}

			return objectMapper.writeValueAsString(schema);

		}
		catch (Exception e) {
			logger.error("Error converting inputSchema to parameters for tool: {}", getName(), e);
			return "{\"type\":\"object\",\"properties\":{},\"additionalProperties\":false}";
		}
	}

	@Override
	public Class<Map<String, Object>> getInputType() {
		@SuppressWarnings("unchecked")
		Class<Map<String, Object>> mapClass = (Class<Map<String, Object>>) (Class<?>) Map.class;
		return mapClass;
	}

	/**
	 * Asynchronous execution - returns CompletableFuture to avoid blocking. This is the
	 * preferred method for non-blocking execution in parallel scenarios.
	 * @param input Tool input parameters
	 * @param toolContext Tool execution context
	 * @return CompletableFuture that completes with the tool execution result
	 */
	@Override
	public CompletableFuture<ToolExecuteResult> applyAsync(Map<String, Object> input, ToolContext toolContext) {
		// Extract toolCallId from ToolContext
		String toolCallId = extractToolCallIdFromContext(toolContext);
		if (toolCallId != null) {
			logger.info("Using provided toolCallId from context: {} for tool: {} (async mode)", toolCallId, getName());

			// Extract planDepth from ToolContext and increment by 1 for subplan
			int parentPlanDepth = extractPlanDepthFromContext(toolContext);
			int subplanDepth = parentPlanDepth + 1;
			logger.info("Parent plan depth: {}, subplan will have depth: {} (async)", parentPlanDepth, subplanDepth);

			return executeSubplanWithToolCallIdAsync(input, toolCallId, subplanDepth);
		}
		else {
			return CompletableFuture.completedFuture(
					new ToolExecuteResult("Error: ToolCallId is required for coordinator tool: " + getName()));
		}
	}

	/**
	 * Synchronous execution - delegates to async version for compatibility. Note: This
	 * will block until the subplan completes.
	 * @param input Tool input parameters
	 * @param toolContext Tool execution context
	 * @return Tool execution result (blocks until async operation completes)
	 */
	@Override
	public ToolExecuteResult apply(Map<String, Object> input, ToolContext toolContext) {
		// Delegate to async version and wait for result
		return applyAsync(input, toolContext).join();
	}

	@Override
	public ToolExecuteResult run(Map<String, Object> input) {
		throw new IllegalArgumentException("ToolCallId is required for coordinator tool: " + getName());
	}

	@Override
	public void cleanup(String planId) {
		// Cleanup logic for the coordinator tool
		logger.debug("Cleaning up coordinator tool: {} for planId: {}", getName(), planId);
	}

	@Override
	public String getCurrentToolStateString() {
		return "Ready";
	}

	// Getter for the wrapped coordinator tool config
	public PlanTemplateConfigVO getCoordinatorToolConfig() {
		return coordinatorToolConfig;
	}

	/**
	 * Extract toolCallId from ToolContext. This method looks for toolCallId in the tool
	 * context that was set by DynamicAgent.
	 * @param toolContext The tool context containing toolCallId information
	 * @return toolCallId if found, null otherwise
	 */
	private String extractToolCallIdFromContext(ToolContext toolContext) {
		try {
			return String.valueOf(toolContext.getContext().get("toolcallId"));

		}
		catch (Exception e) {
			logger.warn("Error extracting toolCallId from context: {}", e.getMessage());
			return null;
		}
	}

	/**
	 * Extract planDepth from ToolContext. This method looks for planDepth in the tool
	 * context that was set by DynamicAgent.
	 * @param toolContext The tool context containing planDepth information
	 * @return planDepth if found, 0 otherwise
	 */
	private int extractPlanDepthFromContext(ToolContext toolContext) {
		try {
			Object depthObj = toolContext.getContext().get("planDepth");
			if (depthObj instanceof Number) {
				return ((Number) depthObj).intValue();
			}
			else if (depthObj instanceof String) {
				return Integer.parseInt((String) depthObj);
			}
			return 0;
		}
		catch (Exception e) {
			logger.warn("Error extracting planDepth from context: {}, defaulting to 0", e.getMessage());
			return 0;
		}
	}

	/**
	 * Execute subplan asynchronously with the provided toolCallId. This method returns a
	 * CompletableFuture to avoid blocking.
	 * @param input The input parameters for the subplan
	 * @param toolCallId The toolCallId to use for this execution
	 * @param planDepth The depth of the subplan in the execution hierarchy
	 * @return CompletableFuture that completes with the tool execution result
	 */
	private CompletableFuture<ToolExecuteResult> executeSubplanWithToolCallIdAsync(Map<String, Object> input,
			String toolCallId, int planDepth) {
		try {
			String planTemplateId = coordinatorToolConfig != null ? coordinatorToolConfig.getPlanTemplateId() : null;
			if (planTemplateId == null) {
				String errorMsg = "Plan template ID is null in coordinator tool config";
				logger.error(errorMsg);
				return CompletableFuture.completedFuture(new ToolExecuteResult(errorMsg));
			}

			logger.info("Executing coordinator tool (async): {} with template: {} and toolCallId: {}", getName(),
					planTemplateId, toolCallId);

			// Get the plan template from PlanTemplateService
			String planJson = planTemplateService.getLatestPlanVersion(planTemplateId);
			if (planJson == null) {
				String errorMsg = "Plan template not found: " + planTemplateId;
				logger.error(errorMsg);
				return CompletableFuture.completedFuture(new ToolExecuteResult(errorMsg));
			}

			// Execute the plan using PlanningCoordinator
			// Generate a new plan ID for this subplan execution using PlanIdDispatcher
			String newPlanId = planIdDispatcher.generateSubPlanId(rootPlanId);

			// Prepare parameters for replacement - add planId to input parameters
			Map<String, Object> parametersForReplacement = new HashMap<>();
			if (input != null) {
				parametersForReplacement.putAll(input);
			}
			// Add the generated planId to parameters
			parametersForReplacement.put("planId", newPlanId);

			// Replace parameter placeholders (<< >>) with actual input parameters
			if (!parametersForReplacement.isEmpty()) {
				try {
					logger.info("Replacing parameter placeholders in plan template with input parameters: {}",
							parametersForReplacement.keySet());
					planJson = parameterMappingService.replaceParametersInJson(planJson, parametersForReplacement);
					logger.debug("Parameter replacement completed successfully");
				}
				catch (Exception e) {
					String errorMsg = "Failed to replace parameters in plan template: " + e.getMessage();
					logger.error(errorMsg, e);
					return CompletableFuture.completedFuture(new ToolExecuteResult(errorMsg));
				}
			}
			else {
				logger.debug("No parameter replacement needed - input: {}", input != null ? input.size() : 0);
			}

			// Parse the JSON to create a PlanInterface
			PlanInterface plan = objectMapper.readValue(planJson, PlanInterface.class);

			// Use the provided toolCallId instead of generating a new one
			logger.info("Using provided toolCallId: {} for subplan execution: {} at depth: {} (async)", toolCallId,
					newPlanId, planDepth);

			// Sub-plans should use the same conversationId as parent, but it's not
			// available in this context
			// Use HTTP_REQUEST as subplans are internal calls
			CompletableFuture<PlanExecutionResult> future = planningCoordinator.executeByPlan(plan, rootPlanId,
					currentPlanId, newPlanId, toolCallId, RequestSource.HTTP_REQUEST, null, planDepth, null);

			// Transform the future result into ToolExecuteResult without blocking
			return future.thenApply(result -> {
				if (result.isSuccess()) {
					String output = result.getFinalResult();
					if (output == null || output.trim().isEmpty()) {
						output = "Subplan executed successfully but no output was generated";
					}
					logger.info("Subplan execution completed successfully (async): {}", output);
					return new ToolExecuteResult(output);
				}
				else {
					String errorMsg = result.getErrorMessage() != null ? result.getErrorMessage()
							: "Subplan execution failed";
					logger.error("Subplan execution failed (async): {}", errorMsg);
					return new ToolExecuteResult("Subplan execution failed: " + errorMsg);
				}
			}).exceptionally(e -> {
				String errorMsg = "Coordinator tool execution failed with exception: " + e.getMessage();
				logger.error("{} for tool: {}", errorMsg, getName(), e);
				return new ToolExecuteResult(errorMsg);
			});

		}
		catch (Exception e) {
			String errorMsg = "Error preparing subplan execution: " + e.getMessage();
			logger.error("{} for tool: {}", errorMsg, getName(), e);
			return CompletableFuture.completedFuture(new ToolExecuteResult(errorMsg));
		}
	}

	@Override
	public boolean isSelectable() {
		// Only selectable if enableInternalToolcall is true
		if (coordinatorToolConfig != null && coordinatorToolConfig.getToolConfig() != null) {
			Boolean enableInternalToolcall = coordinatorToolConfig.getToolConfig().getEnableInternalToolcall();
			return enableInternalToolcall != null && enableInternalToolcall;
		}
		return false;
	}

}
